package com.os.ops;

import com.os.config.RsConfig;
import com.os.data.DataStore;
import com.os.feign.DataService;
import com.os.util.IoUtilPlus;
import lombok.SneakyThrows;
import org.apache.commons.math3.linear.*;
import org.apache.commons.math3.util.ArithmeticUtils;

import java.io.IOException;
import java.io.OutputStream;
import java.net.URI;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
/**
 * @author zhaozhikai
 * @date 2023/12/26 21:50
 */
public class RsPutStream extends OutputStream {
    private TempPutStream[] write;
    //分片缓存大小 = 预读取缓存大小 * 数据分片数量 , 这里的预读取缓存大小可以是 IoUtilPlus的 bufferSize
    private int cacheSize;
    private List<Byte> cache;
    private RsConfig rsConfig;


    public RsPutStream(RsConfig rsConfig, DataService dataService, String hash, int size) {
        this.rsConfig = rsConfig;
        this.cacheSize = rsConfig.DATA_SHARDS * IoUtilPlus.bufferSize;
        this.cache = new ArrayList<>(cacheSize);
        List<URI> uris = DataStore.chooseRandomDataServer(rsConfig.TOTAL_SHARDS, new ArrayList<>());

        //每个分片大小
        int storedSize = size + rsConfig.BYTES_IN_INT;
        int shardSize = (storedSize + rsConfig.DATA_SHARDS - 1) / rsConfig.DATA_SHARDS;

        write = new TempPutStream[uris.size()];
        for(int i = 0; i < uris.size(); i ++){
            write[i] = new TempPutStream(uris.get(i),dataService,String.format("%s.%d",hash, i),shardSize);
        }

    }

    @Override
    public void write(byte[] b) throws IOException {
        write(b, 0, b.length);
    }

    @Override
    public void write(byte[] b, int off, int len) throws IOException {
        int length = len;
        int idx = 0;
        while(length > 0){
            cache.add(b[idx]);
            if(cache.size() == cacheSize){
                flush();
            }
            idx++;
            length--;
        }

    }

    @Override
    public void flush() throws IOException {
        byte[] c = new byte[cache.size()];
        for(int i = 0; i < cache.size(); i ++){
            c[i] = cache.get(i);
        }
        byte[][] bytes = buildShared(c);
        for(int i = 0; i < rsConfig.TOTAL_SHARDS; i ++){
            write[i].write(bytes[i]);
        }
        this.cache.clear();
    }

    @Override
    public void write(int b) throws IOException {
        write(new byte[]{(byte) b}, 0,1);
    }

    public void commit(boolean good) throws IOException {
        flush();
        for(TempPutStream w : write){
            w.commit(good);
        }
    }

    /**
     * 构建数据分片
     *
     * @param file
     * @return
     */
    @SneakyThrows
    @Deprecated
    private byte[][] buildShared(byte[] file) {
        int storedSize = file.length + rsConfig.BYTES_IN_INT;
        int shardSize = (storedSize + rsConfig.DATA_SHARDS - 1) / rsConfig.DATA_SHARDS;
        int fileSize = file.length;
        int bufferSize = shardSize * rsConfig.DATA_SHARDS;
        byte[] allBytes = new byte[bufferSize];
        ByteBuffer.wrap(allBytes).putInt(fileSize);
        System.arraycopy(file,0,  allBytes, rsConfig.BYTES_IN_INT, fileSize);
        byte[][] shards = new byte[rsConfig.TOTAL_SHARDS][shardSize];
        for (int i = 0; i < rsConfig.DATA_SHARDS; i++) {
            System.arraycopy(allBytes, i * shardSize, shards[i], 0, shardSize);
        }
        rsConfig.reedSolomon.encodeParity(shards, 0, shardSize);
        return shards;
    }


}
